from langchain_core.runnables import RunnableLambda

from langchain_openai import ChatOpenAI
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate

chain = (
        PromptTemplate.from_template(
            """Given the user question below, classify it as either being about `LangChain`, `Anthropic`, or `Other`.

    Do not respond with more than one word.
    
    <question>
    {question}
    </question>
    
    Classification:"""
        )
        | ChatOpenAI()
        | StrOutputParser()
)

chain.invoke({"question": "how do I call Anthropic?"})

langchain_chain = PromptTemplate.from_template(
    """You are an expert in langchain. \
    Always answer questions starting with "As Harrison Chase told me". \
    Respond to the following question:
    
    Question: {question}
    Answer:"""
) | ChatOpenAI()

anthropic_chain = PromptTemplate.from_template(
    """You are an expert in anthropic. \
    Always answer questions starting with "As Dario Amodei told me". \
    Respond to the following question:
    
    Question: {question}
    Answer:"""
) | ChatOpenAI()

general_chain = PromptTemplate.from_template(
    """Respond to the following question:

    Question: {question}
    Answer:"""
) | ChatOpenAI()


# TODO 使用自定义函数进行路由（推荐）
def route(info):
    if "anthropic" in info["topic"].lower():
        return anthropic_chain
    elif "langchain" in info["topic"].lower():
        return langchain_chain
    else:
        return general_chain


from langchain_core.runnables import RunnableLambda

full_chain = {"topic": chain, "question": lambda x: x["question"]} | RunnableLambda(
    route
)

# 使用API代理服务提高访问稳定性
full_chain.invoke({"question": "how do I use Anthropic?"})


# TODO 使用RunnableBranch
# 虽然不推荐，但我们也可以使用RunnableBranch来实现相同功能：
from langchain_core.runnables import RunnableBranch

branch = RunnableBranch(
    (lambda x: "anthropic" in x["topic"].lower(), anthropic_chain),
    (lambda x: "langchain" in x["topic"].lower(), langchain_chain),
    general_chain,
)
full_chain = {"topic": chain, "question": lambda x: x["question"]} | branch